from wrf import (getvar, interplevel, smooth2d, to_np, latlon_coords, get_cartopy,
                 cartopy_xlim, cartopy_ylim)
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
from Func_List_Files import list_ncfiles, list_csv_files_3
from cartopy.feature import NaturalEarthFeature
from Func_Extract_Data import Extract_CSV_Data
from matplotlib.cm import get_cmap
import matplotlib.pyplot as plt
import scipy.ndimage as ndimage
from netCDF4 import Dataset
import cartopy.crs as crs
import numpy as np
import math
import csv 
import os

#---------------------------------------------------------------------------------------

plt.rcParams.update({
    "text.usetex": True,
    "font.family": "sans-serif",
    "font.sans-serif": ["Helvetica"]})
# for Palatino and other serif fonts use:
plt.rcParams.update({
    "text.usetex": True,
    "font.family": "serif",
    "font.serif": ["Palatino"],
})

#---------------------------------------------------------------------------------------

Input_Dir = '.../Post_Processing/Snapshot/' # Provide the input directory
HNS       = ['August_19']
HR        = ['168']
GSS       = ['32km']
TMS       = ['NoTurb']
CLZS      = ['CLZ100p0', 'Default', 'CLZ0p0001']
CLZS_Text = ['$\mathit{C}$'+r'$_{Lz = 100.0}$', '$\mathit{Default}$','$\mathit{C}$'+r'$_{Lz = 0.0001}$']
Extents   = [[-81, -11, 0, 50.1]]
Time_Idx  = 0
fig, ax   = plt.subplots(1, 3, figsize=(16,8), sharex=False, sharey=False, subplot_kw={'projection': crs.PlateCarree()}, gridspec_kw={'hspace': 0.01, 'wspace': 0.2})

#---------------------------------------------------------------------------------------
idx = 0
for HN in HNS:
    for GS in GSS:
        for TM in TMS:
            clz_count = 0
            for CLZ in CLZS:
                csv_files, lat, lon, U = [], [], [], []
                Hurricane_Setting      = HN + '_' + HR[idx] + 'Hours_' + GS + '_' + TM + '_' + CLZ
                csv_file               = list_csv_files_3 (Input_Dir + Hurricane_Setting, csv_files)[Time_Idx]
                os.chdir(Input_Dir + Hurricane_Setting)
                with open(csv_file) as f:
                    reader = csv.reader(f)
                    row_header = next(reader)
                    lon = row_header[1:-1]
                    for i in range (len(lon)):
                        lon[i] = float(lon[i])
                    Extract_CSV_Data (Input_Dir + Hurricane_Setting, csv_file, lat, str(row_header[0])) 
                    P   = np.zeros(([int(len(lat)) , int(len(lon))]))

                    j = 0
                    for i in range(len(lon)):
                        Profile = [] 
                        Extract_CSV_Data (Input_Dir + Hurricane_Setting, csv_file, Profile, str(lon[i])) 
                        P[:,j]  = Profile
                        j+=1
#---------------------------------------------------------------------------------------
                states = NaturalEarthFeature(category="cultural", scale="50m",
                                                facecolor="none",
                                                name="admin_1_states_provinces_shp")
                ax[clz_count].add_feature(states, linewidth=.5, edgecolor="black")
                ax[clz_count].coastlines('50m', linewidth=0.8)
                gl = ax[clz_count].gridlines(crs=crs.PlateCarree(), draw_labels=True,
                    linewidth=0, color='black', alpha=0.2, linestyle='--')
                gl.xlabel_style= {'size': 14, 'color': 'black'}
                gl.ylabel_style= {'size': 14, 'color': 'black'}
                gl.top_labels = False
                gl.right_labels = False
                ax[clz_count].set_extent(Extents[idx])

                im = ax[clz_count].contourf(to_np(lon), to_np(lat), to_np(P), 150,
                    transform=crs.PlateCarree(),vmin=-15 ,vmax=15,
                    cmap= plt.cm.RdBu_r, )
                ax[clz_count].set_title(HN +' ' +  CLZS_Text[clz_count], {'size': 16}, pad= 0.3, y = 1.05)
                clz_count += 1                                
    idx += 1
cbar = fig.colorbar(im, ax = ax[:], location = 'right',  pad=0.05, extend='neither')
cbar.set_label('Azimuthal wind velocity \n at the lowest level' + " [m s" + r"$^{-1}$" + ']', rotation = 270, labelpad = 40, size = 'xx-large')
cbar.ax.tick_params(labelsize=14)

plt.show()
